import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics
import os

VER = 'v1.0'
BASE_DIR, FILE_NAME = os.path.split(__file__)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)


class ConvCell(keras.Model):

    def __init__(self, filters, ksize=(3, 3), strides=(1, 1), padding='same', **kwargs):
        super().__init__(**kwargs)
        self.conv = layers.Conv2D(filters, ksize, strides, padding, use_bias=False)
        self.bn = layers.BatchNormalization()
        self.relu = layers.ReLU()

    def call(self, inputs, training=None, mask=None):
        x = self.conv(inputs, training=training)
        x = self.bn(x, training=training)
        x = self.relu(x, training=training)
        return x


class InceptionNetBlock(keras.Model):

    def __init__(self, ch, is_shrink=False, **kwargs):
        super().__init__(**kwargs)
        if is_shrink:
            strides = (2, 2)
        else:
            strides = (1, 1)
        self.layer1 = ConvCell(ch, (3, 3), strides)
        self.layer2_1 = ConvCell(ch, (1, 1), strides)
        self.layer2_2 = ConvCell(ch)
        self.layer3_1 = ConvCell(ch, (1, 1), strides)
        self.layer3_2 = ConvCell(ch)
        self.layer3_3 = ConvCell(ch)
        self.layer4_1 = layers.AvgPool2D((3, 3), (1, 1), 'same')
        self.layer4_2 = ConvCell(ch, (1, 1), strides)

    def call(self, inputs, training=None, mask=None):
        x1 = self.layer1(inputs, training=training)
        x2 = self.layer2_1(inputs, training=training)
        x2 = self.layer2_2(x2, training=training)
        x3 = self.layer3_1(inputs, training=training)
        x3 = self.layer3_2(x3, training=training)
        x4 = self.layer4_1(inputs, training=training)
        x4 = self.layer4_2(x4, training=training)
        x = tf.concat([x1, x2, x3, x4], axis=3)
        return x


class InceptionNet10(keras.Model):

    def __init__(self, ch=32, n_blocks=2, n_cls=1000, **kwargs):
        super().__init__(**kwargs)

        lists = []
        for i in range(n_blocks):
            if i != 0:
                ch *= 2
            for id_layers in range(2):
                if 0 == id_layers:
                    shrink = True
                else:
                    shrink = False
                lists.append(InceptionNetBlock(ch, shrink))
        self.convs = keras.Sequential(lists)
        self.fc = layers.Dense(n_cls)

    def call(self, inputs, training=None, mask=None):
        x = self.convs(inputs, training=training)
        x = layers.GlobalAvgPool2D()(x)
        x = self.fc(x)
        return x


model = InceptionNet10(32, 2, 10)
model.build(input_shape=(None, 32, 32, 3))
model.summary()

x = tf.zeros((4, 32, 32, 3), dtype=tf.float32)
pred = model(x)
print('pred:', tf.shape(pred))
